import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class gat_layer(nn.Module):
    def __init__(self, insize, outsize):
        super(gat_layer, self).__init__()
        self.insize = insize
        self.outsize = outsize
        self.weight = nn.Parameter(torch.FloatTensor(insize, outsize))
        self.a = nn.Parameter(torch.FloatTensor(2 * outsize, 1))
        self.reset_parameters()
        
    # 初始化参数
    def reset_parameters(self):
        stdv = 1. / math.sqrt(self.weight.size(1))
        self.weight.data.uniform_(-stdv, stdv)
        self.a.data.uniform_(-stdv, stdv)
        # self.weight.data.fill_(0)
        # self.a.data.fill_(0)

	# 前向传播函数
    def forward(self, hl, adj):
        mid = torch.mm(hl, self.weight)
        N = hl.size()[0]    # 图的节点数
        
        # 得到任意两个点属性两两拼接的三维矩阵，第一维代表拼接的前半段，第二维代表后半段，第三维是拼接后属性
        concat = torch.cat([mid.repeat(1, N).view(N * N, -1), mid.repeat(N, 1)], dim=1).view(N, -1, 2 * self.outsize)
        # 得到a与每个拼接属性的乘积矩阵并激活
        all_attention = F.leaky_relu(torch.matmul(concat, self.a).squeeze(2))

        zero = -1e20 * torch.ones_like(all_attention)
        attention = torch.where(adj>0, all_attention, zero)
        attention = F.softmax(attention, dim=1)
        hn = torch.matmul(attention, mid)

        return hn


class gat(nn.Module):
    def __init__(self, insize, outsize, hidsize, hidlayernum):
        super(gat, self).__init__()
        self.ly1 = gat_layer(insize, hidsize)
        self.ly2 = gat_layer(hidsize, outsize)
        self.hidlayernum = hidlayernum
        if hidlayernum < 0:
            raise ValueError("hidlayernum < 0")
        else:
            self.hid = []
            for i in range(hidlayernum):
                self.hid.append(gat_layer(hidsize, hidsize))


    def forward(self, feature, adj):
        hid_out = F.relu(self.ly1(feature, adj))
        for i in range(self.hidlayernum):
            hid_out = F.relu(self.hid[i](hid_out, adj))
        ly2_out = self.ly2(hid_out, adj)
        res = F.log_softmax(ly2_out, dim=1)
        return res

